{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "4a87f55f-9cb5-4aee-8d80-cd3b7935aa2d",
   "metadata": {},
   "source": [
    "# Strands multi-agent System with AgentCore Memory Tool (Short term Memory)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e1661265-84cd-4a1a-9a62-739c925af6ca",
   "metadata": {},
   "source": [
    "## Introduction\n",
    "\n",
    "This notebook demonstrates how to implement a **multi-agent system with shared memory** using AWS AgentCore Memory and the Strands framework. While our previous examples focused on single-agent memory, this notebook explores how multiple specialized agents can work together while accessing a common memory store.\n",
    "\n",
    "## Tutorial Details\n",
    "\n",
    "| Information         | Details                                                                          |\n",
    "|:--------------------|:---------------------------------------------------------------------------------|\n",
    "| Tutorial type       | Short Term Conversational                                                        |\n",
    "| Agent usecase       | Travel Planning Assistant                                                        |\n",
    "| Agentic Framework   | Strands Agents                                                                   |\n",
    "| LLM model           | Anthropic Claude Sonnet 3                                                        |\n",
    "| Tutorial components | AgentCore Short-term Memory, Strands Agents, Memory retrieval via Tool           |\n",
    "| Example complexity  | Beginner                                                                         |\n",
    "\n",
    "\n",
    "What you will learn:\n",
    "\n",
    "- How to set up a shared memory resource that multiple agents can access\n",
    "- Creating specialized agents as tools with their own memory access\n",
    "- Implementing a coordinator agent that delegates to specialized agents\n",
    "- Maintaining conversation context across multiple agent interactions\n",
    "\n",
    "### Scenario context\n",
    "\n",
    "In this example, we'll create a **Travel Planning System** with:\n",
    "1. A Flight Booking Assistant specialized in air travel\n",
    "2. A Hotel Booking Assistant focused on accommodations\n",
    "3. A Travel Coordinator that delegates to these specialized agents\n",
    "\n",
    "This approach demonstrates how complex domains can be broken down into specialized agents that share memory the same memory store.\n",
    "\n",
    "## Architecture\n",
    "<div style=\"text-align:left\">\n",
    "    <img src=\"architecture.png\" width=\"65%\" />\n",
    "</div>\n",
    "\n",
    "## Prerequisites\n",
    "- Python 3.10+\n",
    "- AWS account with appropriate permissions\n",
    "- AWS IAM role with appropriate permissions for AgentCore Memory\n",
    "- Access to Amazon Bedrock models\n",
    "\n",
    "Let's get started by setting up our environment and creating our shared memory resource!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "091a7ba3-0933-43a2-93b8-3379e73440c1",
   "metadata": {},
   "source": [
    "## Step 1: Environment set up\n",
    "Let's begin importing all the necessary libraries and defining the clients to make this notebook work."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ea35f65-29ba-42e7-a756-7ab5b93a84ce",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -qr requirements.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f63a5236-2e64-4029-876c-87fe808ce0a8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import logging\n",
    "from datetime import datetime\n",
    "from strands.hooks import AgentInitializedEvent, HookProvider, HookRegistry, MessageAddedEvent"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bff35a0f-79bc-4d47-81d4-5a52f219e664",
   "metadata": {},
   "source": [
    "Define the region and the role with the appropiate permissions for Amazon Bedrock models and AgentCore"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "746a9e5c-0893-431a-9585-9e8bfb13fd52",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "region = os.getenv('AWS_REGION', 'us-west-2')\n",
    "\n",
    "logging.basicConfig(level=logging.INFO, format=\"%(asctime)s - %(levelname)s - %(message)s\", datefmt=\"%Y-%m-%d %H:%M:%S\")\n",
    "logger = logging.getLogger(\"agentcore-memory\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ef644e57-9e6f-413d-9b26-5eeb775d7760",
   "metadata": {},
   "source": [
    "## Step 2: Creating Shared Memory\n",
    "In this section, we'll create a memory resource that will be shared among our specialized agents."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3be50b1e-e11b-43e4-8ba6-3c9d703110eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "from bedrock_agentcore.memory import MemoryClient"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "99f7fd67-2e2d-4248-8353-b0c495f9fd93",
   "metadata": {},
   "outputs": [],
   "source": [
    "client = MemoryClient(region_name=region)\n",
    "memory_name = \"TravelAgent_STM_%s\" % datetime.now().strftime(\"%Y%m%d%H%M%S\")\n",
    "memory_id = None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8f32d7a9-dcdd-4296-968c-94aa7f3a7c30",
   "metadata": {},
   "outputs": [],
   "source": [
    "from botocore.exceptions import ClientError\n",
    "\n",
    "try:\n",
    "    print(\"Creating Memory...\")\n",
    "    memory_name = memory_name\n",
    "\n",
    "    # Create the memory resource\n",
    "    memory = client.create_memory_and_wait(\n",
    "        name=memory_name,                       # Unique name for this memory store\n",
    "        description=\"Travel Agent STM\",         # Human-readable description\n",
    "        strategies=[],                          # No special memory strategies for short-term memory\n",
    "        event_expiry_days=7,                    # Memories expire after 7 days\n",
    "        max_wait=300,                           # Maximum time to wait for memory creation (5 minutes)\n",
    "        poll_interval=10                        # Check status every 10 seconds\n",
    "    )\n",
    "\n",
    "    # Extract and print the memory ID\n",
    "    memory_id = memory['memoryId']\n",
    "    print(f\"Memory created successfully with ID: {memory_id}\")\n",
    "except ClientError as e:\n",
    "    if e.response['Error']['Code'] == 'ValidationException' and \"already exists\" in str(e):\n",
    "        # If memory already exists, retrieve its ID\n",
    "        memories = client.list_memories()\n",
    "        memory_id = next((m['id'] for m in memories if m['id'].startswith(memory_name)), None)\n",
    "        logger.info(f\"Memory already exists. Using existing memory ID: {memory_id}\")\n",
    "except Exception as e:\n",
    "    # Handle any errors during memory creation\n",
    "    print(f\"❌ ERROR: {e}\")\n",
    "    import traceback\n",
    "    traceback.print_exc()\n",
    "\n",
    "    # Cleanup on error - delete the memory if it was partially created\n",
    "    if memory_id:\n",
    "        try:\n",
    "            client.delete_memory_and_wait(memory_id=memory_id)\n",
    "            logger.info(f\"Cleaned up memory: {memory_id}\")\n",
    "        except Exception as cleanup_error:\n",
    "            logger.info(f\"Failed to clean up memory: {cleanup_error}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aa3fed1c-aa71-409b-9b37-2ee03b0aeeb5",
   "metadata": {},
   "source": [
    "### Understanding Shared Memory for Multi-Agent Systems\n",
    "\n",
    "The memory resource we've created will serve as a shared knowledge base for our travel planning system. All agents will read from and write to this common memory store, enabling:\n",
    "\n",
    "1. **Knowledge Consistency**: All agents work with the same information\n",
    "2. **Context Preservation**: Conversation history is maintained across agent transitions\n",
    "3. **Specialized Access**: Each agent will have its own actor_id but share the session_id\n",
    "\n",
    "This approach allows specialized agents to focus on their domains while still benefiting from the full conversation context."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5287c609-2117-41db-8ab4-5df5acdbf150",
   "metadata": {},
   "source": [
    "## Step 3: Create Memory Hook Provider\n",
    "\n",
    "This step defines our custom `MemoryHookProvider` class that automates memory operations. Hooks are special functions that run at specific points in an agent's execution lifecycle. The memory hook we're creating serves two primary functions:\n",
    "\n",
    "1. **Retrieve Memories**: Automatically fetches relevant past conversations when a user sends a message\n",
    "2. **Save Memories**: Stores new conversations after the agent responds\n",
    "\n",
    "This creates a seamless memory experience without manual management."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "371823d3-186c-4c99-ab78-9316e522afd3",
   "metadata": {},
   "outputs": [],
   "source": [
    "class ShortTermMemoryHook(HookProvider):\n",
    "    def __init__(self, memory_client: MemoryClient, memory_id: str, actor_id: str, session_id: str):\n",
    "        self.memory_client = memory_client\n",
    "        self.memory_id = memory_id\n",
    "        self.actor_id = actor_id\n",
    "        self.session_id = session_id\n",
    "    \n",
    "    def on_agent_initialized(self, event: AgentInitializedEvent):\n",
    "        \"\"\"Load recent conversation history when agent starts\"\"\"\n",
    "        try:\n",
    "            # Get last 5 conversation turns\n",
    "            recent_turns = self.memory_client.get_last_k_turns(\n",
    "                memory_id=self.memory_id,\n",
    "                actor_id=self.actor_id,\n",
    "                session_id=self.session_id,\n",
    "                k=5,\n",
    "                branch_name=\"main\"\n",
    "            )\n",
    "            \n",
    "            if recent_turns:\n",
    "                # Format conversation history for context\n",
    "                context_messages = []\n",
    "                for turn in recent_turns:\n",
    "                    for message in turn:\n",
    "                        role = message['role'].lower()\n",
    "                        content = message['content']['text']\n",
    "                        context_messages.append(f\"{role.title()}: {content}\")\n",
    "                \n",
    "                context = \"\\n\".join(context_messages)\n",
    "                logger.info(f\"Context from memory: {context}\")\n",
    "                \n",
    "                # Add context to agent's system prompt\n",
    "                event.agent.system_prompt += f\"\\n\\nRecent conversation history:\\n{context}\\n\\nContinue the conversation naturally based on this context.\"\n",
    "                \n",
    "                logger.info(f\"✅ Loaded {len(recent_turns)} recent conversation turns\")\n",
    "            else:\n",
    "                logger.info(\"No previous conversation history found\")\n",
    "                \n",
    "        except Exception as e:\n",
    "            logger.error(f\"Failed to load conversation history: {e}\")\n",
    "    \n",
    "    def on_message_added(self, event: MessageAddedEvent):\n",
    "        \"\"\"Store conversation turns in memory\"\"\"\n",
    "        messages = event.agent.messages\n",
    "        try:\n",
    "            self.memory_client.create_event(\n",
    "                memory_id=self.memory_id,\n",
    "                actor_id=self.actor_id,\n",
    "                session_id=self.session_id,\n",
    "                messages=[(messages[-1][\"content\"][0][\"text\"], messages[-1][\"role\"])]\n",
    "            )\n",
    "            \n",
    "        except Exception as e:\n",
    "            logger.error(f\"Failed to store message: {e}\")\n",
    "    \n",
    "    def register_hooks(self, registry: HookRegistry) -> None:\n",
    "        # Register memory hooks\n",
    "        registry.add_callback(MessageAddedEvent, self.on_message_added)\n",
    "        registry.add_callback(AgentInitializedEvent, self.on_agent_initialized)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11313704-f182-423f-96d7-a64b4482c8aa",
   "metadata": {},
   "source": [
    "## Step 4: Create Multi-Agent Architecture with Strands Agents\n",
    "In this section, we'll create our multi-agent system with specialized agents for flight and hotel bookings, both sharing access to our memory resource."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a288afca-e06e-48b2-b10b-9cb0465c9e23",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import the necessary components\n",
    "from strands import Agent, tool"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b1293c57-967a-4cba-ad18-a07e97358929",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create unique actor IDs for each specialized agent but share the session ID\n",
    "flight_actor_id = f\"flight-user-{datetime.now().strftime('%Y%m%d%H%M%S')}\"\n",
    "hotel_actor_id = f\"hotel-user-{datetime.now().strftime('%Y%m%d%H%M%S')}\"\n",
    "session_id = f\"travel-session-{datetime.now().strftime('%Y%m%d%H%M%S')}\"\n",
    "flight_namespace = f\"travel/{flight_actor_id}/preferences\"\n",
    "hotel_namespace = f\"travel/{hotel_actor_id}/preferences\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "48f39b9d-4238-45e5-a8eb-98a8c450e65a",
   "metadata": {},
   "source": [
    "### Creating Specialized Agents with Memory Access\n",
    "\n",
    "Next, we'll define system prompts for our specialized agents. Each prompt includes the memory parameters in a format that the agent can parse:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bb748408-80cc-4f2a-8950-ed1523bec2aa",
   "metadata": {},
   "outputs": [],
   "source": [
    "# System prompt for the hotel booking specialist\n",
    "HOTEL_BOOKING_PROMPT = f\"\"\"You are a hotel booking assistant. Help customers find hotels, make reservations, and answer questions about accommodations and amenities. \n",
    "Provide clear information about availability, pricing, and booking procedures in a friendly, helpful manner.\"\"\"\n",
    "\n",
    "# System prompt for the flight booking specialist\n",
    "FLIGHT_BOOKING_PROMPT = f\"\"\"You are a flight booking assistant. Help customers find flights, make reservations, and answer questions about airlines, routes, and travel policies. \n",
    "Provide clear information about flight availability, pricing, schedules, and booking procedures in a friendly, helpful manner.\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a42e6a49-dc31-48ae-b6ca-5e08266fd74a",
   "metadata": {},
   "source": [
    "### Implementing Agent Tools\n",
    "Now we'll implement our specialized agents as tools that can be used by the coordinator agent:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f67ed8c-c15b-4e5a-a130-aa3ef165fc3f",
   "metadata": {},
   "outputs": [],
   "source": [
    "@tool\n",
    "def flight_booking_assistant(query: str) -> str:\n",
    "    \"\"\"\n",
    "    Process and respond to flight booking queries.\n",
    "\n",
    "    Args:\n",
    "        query: A flight-related question about bookings, schedules, airlines, or travel policies\n",
    "\n",
    "    Returns:\n",
    "        Detailed flight information, booking options, or travel advice\n",
    "    \"\"\"\n",
    "    try:\n",
    "        flight_memory_hooks = ShortTermMemoryHook(\n",
    "            memory_id=memory_id,\n",
    "            memory_client=client,\n",
    "            actor_id=flight_actor_id,\n",
    "            session_id=session_id\n",
    "        )\n",
    "        \n",
    "        flight_agent = Agent(hooks=[flight_memory_hooks], system_prompt=FLIGHT_BOOKING_PROMPT)\n",
    "\n",
    "        # Call the agent and return its response\n",
    "        response = flight_agent(query)\n",
    "        return str(response)\n",
    "    except Exception as e:\n",
    "        return f\"Error in flight booking assistant: {str(e)}\"\n",
    "\n",
    "@tool\n",
    "def hotel_booking_assistant(query: str) -> str:\n",
    "    \"\"\"\n",
    "    Process and respond to hotel booking queries.\n",
    "\n",
    "    Args:\n",
    "        query: A hotel-related question about accommodations, amenities, or reservations\n",
    "\n",
    "    Returns:\n",
    "        Detailed hotel information, booking options, or accommodation advice\n",
    "    \"\"\"\n",
    "    try:\n",
    "        hotel_memory_hooks = ShortTermMemoryHook(\n",
    "            memory_id=memory_id,\n",
    "            memory_client=client,\n",
    "            actor_id=hotel_actor_id,\n",
    "            session_id=session_id\n",
    "        )\n",
    "\n",
    "        hotel_booking_agent = Agent(hooks=[hotel_memory_hooks], system_prompt=HOTEL_BOOKING_PROMPT)\n",
    "        \n",
    "        # Call the agent and return its response\n",
    "        response = hotel_booking_agent(query)\n",
    "        return str(response)\n",
    "    except Exception as e:\n",
    "        return f\"Error in hotel booking assistant: {str(e)}\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e7b572d1-d429-4678-a318-2ab153408e23",
   "metadata": {},
   "source": [
    "### Creating the Coordinator Agent\n",
    "\n",
    "Finally, we'll create the main travel planning agent that coordinates between these specialized tools:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4abe237d-36c2-42ab-ae54-4554e1e4f23a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# System prompt for the coordinator agent\n",
    "TRAVEL_AGENT_SYSTEM_PROMPT = \"\"\"\n",
    "You are a comprehensive travel planning assistant that coordinates between specialized tools:\n",
    "- For flight-related queries (bookings, schedules, airlines, routes) → Use the flight_booking_assistant tool\n",
    "- For hotel-related queries (accommodations, amenities, reservations) → Use the hotel_booking_assistant tool\n",
    "- For complete travel packages → Use both tools as needed to provide comprehensive information\n",
    "- For general travel advice or simple travel questions → Answer directly\n",
    "\n",
    "Each agent will have its own memory in case the user asks about historic data.\n",
    "When handling complex travel requests, coordinate information from both tools to create a cohesive travel plan.\n",
    "Provide clear organization when presenting information from multiple sources. \\\n",
    "Ask max two questions per turn. Keep the messages short, don't overwhelm the customer.\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f10bb9f-031e-49b6-8175-1b676cdd2576",
   "metadata": {},
   "outputs": [],
   "source": [
    "travel_agent = Agent(\n",
    "    system_prompt=TRAVEL_AGENT_SYSTEM_PROMPT,\n",
    "    tools=[flight_booking_assistant, hotel_booking_assistant]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c361c59d-e322-47ae-92a8-ee6cd994df22",
   "metadata": {},
   "source": [
    "#### Your Multi-Agent System is ready !!\n",
    "\n",
    "## Let's test the Agent.\n",
    "\n",
    "Let's test our multi-agent system with a travel planning scenario:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8b13aa5d-4abc-41eb-916b-35ccf0b5ab03",
   "metadata": {},
   "outputs": [],
   "source": [
    "response = travel_agent(\"Hello, I would like to book a trip from LA to Madrid. From July 1 to August 2.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec209fa8-fd5a-41a8-b7d1-86a1f137331f",
   "metadata": {},
   "outputs": [],
   "source": [
    "response = travel_agent(\"I would only like to focus on the flight at the moment. direct flimid-range, city center, pool, standard room\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ef1ed3c9-93ee-4617-81a7-b462133f4488",
   "metadata": {},
   "source": [
    "## Testing Memory Persistence\n",
    "\n",
    "To test if our memory system is working correctly, we'll create a new instance of the travel agent and see if it can access the previously stored information:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8007b15a-2b00-4f07-aa7d-dfa8ba7b5b10",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a new instance of the travel agent\n",
    "new_travel_agent = Agent(\n",
    "    system_prompt=TRAVEL_AGENT_SYSTEM_PROMPT,\n",
    "    tools=[flight_booking_assistant, hotel_booking_assistant]\n",
    ")\n",
    "\n",
    "# Ask about previous conversations\n",
    "new_travel_agent(\"Can you remind me about flights talked about before?\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2fb6cdd1-5c11-4346-9c5a-93fb44c779ac",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "In this notebook, we've demonstrated:\n",
    "\n",
    "1. How to create a shared memory resource for multiple agents\n",
    "2. How to implement specialized agents as tools with memory access\n",
    "3. How to coordinate between multiple agents while maintaining conversation context\n",
    "4. How memory persists across different agent instances\n",
    "\n",
    "This multi-agent architecture with shared memory provides a powerful approach for building complex conversational AI systems that can handle specialized domains while maintaining a cohesive user experience."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18852f04-e06e-4787-9c4f-3d2d2ba8563b",
   "metadata": {},
   "source": [
    "## Clean up\n",
    "Let's delete the memory to clean up the resources used in this notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3d4577ee-3a79-4a44-abde-86391e98e4bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "#client.delete_memory_and_wait(\n",
    "#        memory_id = memory_id,\n",
    "#        max_wait = 300,\n",
    "#        poll_interval =10\n",
    "#)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "conda_python3",
   "language": "python",
   "name": "conda_python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
